/-
Copyright (c) 2022 Jannis Limperg. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Jannis Limperg
-/

import Aesop.Check
import Aesop.Frontend.Attribute
import Aesop.Options
import Aesop.RuleSet
import Aesop.Search.Expansion
import Aesop.Search.ExpandSafePrefix
import Aesop.Search.Queue
import Aesop.Tree
import Aesop.Util

open Lean
open Lean.Elab.Tactic (liftMetaTacticAux TacticM)
open Lean.Parser.Tactic (tacticSeq)
open Lean.Meta

namespace Aesop

variable [Aesop.Queue Q]

partial def nextActiveGoal : SearchM Q GoalRef := do
  let some gref ← popGoal?
    | throwError "aesop/expandNextGoal: internal error: no active goals left"
  if ! (← (← gref.get).isActive) then
    nextActiveGoal
  else
    return gref

def expandNextGoal : SearchM Q Unit := do
  let gref ← nextActiveGoal
  let g ← gref.get
  let (initialGoal, initialMetaState) ← g.currentGoalAndMetaState
  let result ← withAesopTraceNode .steps
    (fmt g.id g.priority initialGoal initialMetaState) do
    initialMetaState.runMetaM' do
      aesop_trace[steps] "Initial goal:{indentD initialGoal}"
    let maxRappDepth := (← read).options.maxRuleApplicationDepth
    if maxRappDepth != 0 && (← gref.get).depth >= maxRappDepth then
      aesop_trace[steps] "Treating the goal as unprovable since it is beyond the maximum rule application depth ({maxRappDepth})."
      gref.markForcedUnprovable
      setMaxRuleApplicationDepthReached
      return .failed
    let result ← expandGoal gref
    let currentIteration ← getIteration
    gref.modify λ g => g.setLastExpandedInIteration currentIteration
    if ← (← gref.get).isActive then
      enqueueGoals #[gref]
    return result
  match result with
  | .proved newRapps | .succeeded newRapps => traceNewRapps newRapps
  | .failed => return
  where
    fmt (id : GoalId) (priority : Percent) (initialGoal : MVarId)
        (initialMetaState : Meta.SavedState)
        (result : Except Exception RuleResult) : SearchM Q MessageData := do
      let tgt ← initialMetaState.runMetaM' do
        initialGoal.withContext do
          addMessageContext $ toMessageData (← initialGoal.getType)
      return m!"{exceptRuleResultToEmoji (·.toEmoji) result} (G{id}) [{priority.toHumanString}] ⋯ ⊢ {tgt}"

    traceNewRapps (newRapps : Array RappRef) : SearchM Q Unit := do
      aesop_trace[steps] do
        for rref in newRapps do
          let r ← rref.get
          r.withHeadlineTraceNode .steps
            (transform := λ msg => return m!"{newNodeEmoji} " ++ msg) do
            withAesopTraceNode .steps (λ _ => return "Metadata") do
              r.traceMetadata .steps
          r.metaState.runMetaM' do
            r.forSubgoalsM λ gref => do
              let g ← gref.get
              g.withHeadlineTraceNode .steps
                (transform := λ msg => return m!"{newNodeEmoji} " ++ msg) do
                aesop_trace![steps] g.preNormGoal
                withAesopTraceNode .steps (λ _ => return "Metadata") do
                  g.traceMetadata .steps

def checkGoalLimit : SearchM Q (Option MessageData) := do
  let maxGoals := (← read).options.maxGoals
  let currentGoals := (← getTree).numGoals
  if maxGoals != 0 && currentGoals >= maxGoals then
    return m!"maximum number of goals ({maxGoals}) reached. Set the 'maxGoals' option to increase the limit."
  return none

def checkRappLimit : SearchM Q (Option MessageData) := do
  let maxRapps := (← read).options.maxRuleApplications
  let currentRapps := (← getTree).numRapps
  if maxRapps != 0 && currentRapps >= maxRapps then
    return m!"maximum number of rule applications ({maxRapps}) reached. Set the 'maxRuleApplications' option to increase the limit."
  return none

def checkRootUnprovable : SearchM Q (Option MessageData) := do
  let root := (← getTree).root
  if (← root.get).state.isUnprovable then
    let msg :=
      if ← wasMaxRuleApplicationDepthReached then
        m!"failed to prove the goal. Some goals were not explored because the maximum rule application depth ({(← read).options.maxRuleApplicationDepth}) was reached. Set option 'maxRuleApplicationDepth' to increase the limit."
      else
        m!"failed to prove the goal after exhaustive search."
    return msg
  return none

def getProof? : SearchM Q (Option Expr) := do
  let (some proof) ← getExprMVarAssignment? (← getRootMVarId)
    | return none
  instantiateMVars proof

private def withPPAnalyze [Monad m] [MonadWithOptions m] (x : m α) : m α :=
  withOptions (·.setBool `pp.analyze true) x

def finalizeProof : SearchM Q Unit := do
  (← getRootMVarId).withContext do
    extractProof
    let (some proof) ← getProof? | throwError
      "aesop: internal error: root goal is proven but its metavariable is not assigned"
    if proof.hasExprMVar then
      let inner :=
        m!"Proof: {proof}\nUnassigned metavariables: {(← getMVarsNoDelayed proof).map (·.name)}"
      throwError "aesop: internal error: extracted proof has metavariables.{indentD inner}"
    withPPAnalyze do
      aesop_trace[proof] "Final proof:{indentExpr proof}"

open Lean.Elab.Tactic in
def checkRenderedScript (script : Array Syntax.Tactic) : SearchM Q Unit := do
  let initialState := (← read).originalMetaState
  let rootGoal := (← read).originalGoal
  let go : TacticM Unit := do
    setGoals [rootGoal]
    evalTactic $ ← `(tacticSeq| $script:tactic*)
    unless (← getUnsolvedGoals).isEmpty do
      throwError "script executed successfully but did not solve the main goal"
  try
    discard $ show MetaM _ from withoutModifyingState do
      initialState.restore
      go.run { elaborator := .anonymous, recover := false }
        |>.run { goals := [rootGoal] } |>.run
  catch e => throwError
    "{Check.script.name}: error while executing generated script:{indentD e.toMessageData}"

def traceScript : SearchM Q Unit := do
  let options := (← read).options
  if ! options.generateScript then
    return
  try
    let script ←
      (← getRootMVarCluster).extractScript (← read).preprocessingScript
    let goal := (← read).originalGoal
    let tacticState ← TacticState.ofGoals #[goal]
    let script ← script.toStructuredScript tacticState
    let script ← script.render tacticState
    if options.traceScript then
      let scriptMsg :=
        MessageData.joinSep (script.map toMessageData |>.toList) "\n"
      withPPAnalyze do
        logInfo m!"Try this:{indentD scriptMsg}"
    if ← Check.script.isEnabled then
      checkRenderedScript script
  catch e =>
    logError m!"aesop: error while generating tactic script:{indentD e.toMessageData}"

def traceTree : SearchM Q Unit := do
  (← (← getRootGoal).get).traceTree .tree

def finishIfProven : SearchM Q Bool := do
  unless (← (← getRootMVarCluster).get).state.isProven do
    return false
  finalizeProof
  traceScript
  traceTree
  return true

-- When we hit a non-fatal error (i.e. the search terminates without a proof
-- because the root goal is unprovable or because we hit a search limit), we
-- usually:
--
-- - Expand all safe rules as much as possible, starting from the root node,
--   until we hit an unsafe rule. We call this the safe prefix.
-- - Extract the proof term for the safe prefix and report the remaining goals.
--
-- The first step is necessary because a goal can become unprovable due to a
-- sibling being unprovable, without the goal ever being expanded. So if we did
-- not expand the safe rules after the fact, the tactic's output would be
-- sensitive to minor changes in, e.g., rule priority.
def handleNonfatalError (err : MessageData) : SearchM Q (Array MVarId) := do
  let opts := (← read).options
  if opts.terminal then
    throwTacticEx `aesop (← getRootMVarId) err
  if opts.warnOnNonterminal then
    logWarning m!"aesop: {err}"
  expandSafePrefix
  let goals ← extractSafePrefix
  aesop_trace[proof] do
    match ← getProof? with
    | some proof =>
      (← getRootMVarId).withContext do
        aesop_trace![proof] "{proof}"
    | none => aesop_trace![proof] "<no proof>"
  traceTree
  return goals

def handleFatalError (e : Exception) : SearchM Q α := do
  traceTree
  throw e

partial def searchLoop : SearchM Q (Array MVarId) :=
  withIncRecDepth do
    if let (some err) ← checkRootUnprovable then
      handleNonfatalError err
    else if ← finishIfProven then
      return #[]
    else if let (some err) ← checkGoalLimit then
      handleNonfatalError err
    else if let (some err) ← checkRappLimit then
      handleNonfatalError err
    else
      expandNextGoal
      checkInvariantsIfEnabled
      incrementIteration
      searchLoop

def preprocessGoal (mvarId : MVarId) (mvars : HashSet MVarId)
    (generateScript : Bool) : MetaM (MVarId × UnstructuredScript) := do
  let (mvarId', _, scriptBuilder?) ←
    renameInaccessibleFVarsWithScript mvarId generateScript
  let script ←
    if let some scriptBuilder := scriptBuilder? then
      pure #[{
        tacticSeq := ← scriptBuilder.unstructured.run
        inGoal := mvarId
        outGoals := #[⟨mvarId', mvars⟩]
        otherSolvedGoals := {}
      }]
    else
      pure #[]
  return (mvarId', script)

def search (goal : MVarId) (ruleSet? : Option RuleSet := none)
     (options : Aesop.Options := {}) (simpConfig : Aesop.SimpConfig := {})
     (simpConfigSyntax? : Option Term := none)
     (profile : Profile := {}) :
     MetaM (Array MVarId × Profile) := do
  goal.checkNotAssigned `aesop
  let originalMetaState ← saveState
  let mvars ← goal.getMVarDependencies
  let originalGoal := goal
  let options ← options.toOptions'
  let (goal, preprocessingScript) ←
    preprocessGoal goal mvars options.generateScript
  let ruleSet ←
    match ruleSet? with
    | none =>
        Frontend.getDefaultRuleSet (includeGlobalSimpTheorems := true)
          options.toOptions
    | some ruleSet => pure ruleSet
  let ⟨Q, _⟩ := options.queue
  let (goals, state, _) ←
    SearchM.run ruleSet options simpConfig simpConfigSyntax? goal
        originalMetaState originalGoal preprocessingScript profile do
      show SearchM Q _ from
      try searchLoop
      catch e => handleFatalError e
      finally freeTree
  return (goals, state.profile)

end Aesop
